/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022. All rights reserved.
 * Description: 性能评估器相关头文件
 * Author: qingjicheng
 * Create: 2022-8-1
 */

#ifndef GRAPH_COMMON_DEF_H
#define GRAPH_COMMON_DEF_H

#include <vector>
#include <queue>
#include <set>
#include <unordered_map>
#include <unordered_set>
#include <algorithm>
#include <chrono>
#include <securec.h>

#include <memory>

#include "hccl_evaluator_api.h"
#include "hccl_types.h"

namespace HcclEvaluator {
using s8 = signed char;
using s16 = signed short;
using s32 = signed int;
using s64 = signed long long;
using u8 = unsigned char;
using u16 = unsigned short;
using u32 = unsigned int;
using u64 = unsigned long long;

constexpr u32 ERR_NODE = 0xFFFFFFFF;

template<typename Element>
class Matrix {
public:
    explicit Matrix() = default;
    virtual ~Matrix()
    {
        if (isOwner_) {
            delete[] beginAddr_;
        }
    }

    Matrix(const Matrix<Element>& that)
    {
        beginAddr_ = new (std::nothrow) Element[that.size_ * that.size_]();
        size_ = that.size_;
        isOwner_ = true;
        if (memcpy_s(beginAddr_, that.size_ * that.size_ * sizeof(Element), that.beginAddr_,
                     that.size_ * that.size_ * sizeof(Element)) != EOK) {
            delete[] beginAddr_;
            beginAddr_ = nullptr;
            size_ = 0;
            isOwner_ = false;
        }
    }

    Matrix &operator=(const Matrix<Element> &that)
    {
        beginAddr_ = new (std::nothrow) Element[that.size_ * that.size_]();
        size_ = that.size_;
        isOwner_ = true;
        if (memcpy_s(beginAddr_, that.size_ * that.size_ * sizeof(Element), that.beginAddr_,
                     that.size_ * that.size_ * sizeof(Element)) != EOK) {
            delete[] beginAddr_;
            beginAddr_ = nullptr;
            size_ = 0;
            isOwner_ = false;
        }
        return *this;
    }

    Matrix(Matrix<Element> &&that)
    {
        beginAddr_ = that.beginAddr_;
        size_ = that.size_;
        isOwner_ = that.isOwner_;
        that.isOwner_ = false;
    }

    Matrix &operator=(Matrix<Element> &&that)
    {
        beginAddr_ = that.beginAddr_;
        size_ = that.size_;
        isOwner_ = that.isOwner_;
        that.isOwner_ = false;
        return *this;
    }

    HcclResult Init(u32 size)
    {
        beginAddr_ = new (std::nothrow) Element[size * size]();
        if (beginAddr_ == nullptr) {
            return HCCL_E_PTR;
        }
        size_ = size;
        isOwner_ = true;
        return HCCL_SUCCESS;
    }

    inline Element &operator() (u32 u, u32 v)
    {
        return *(beginAddr_ + u * size_ + v);
    }
private:
    Element *beginAddr_ = nullptr;
    u32 size_;
    bool isOwner_ = false;
};

enum class AdjaMatrixAttr {
    NOT_DETECT,
    RING,
    FULL_MESH,
    STAR,
    MULTI_RING,
    NO_TOPO
};

struct AdjaMatrixEdge {
    HcclPort *ports = nullptr;
    inline bool IsConnected()
    {
        return ports != nullptr;
    }
};

template<typename Edge = AdjaMatrixEdge, typename Attr = AdjaMatrixAttr>
class AdjacencyMatrix {
public:
    explicit AdjacencyMatrix() = default;
    AdjacencyMatrix(AdjacencyMatrix<Edge, Attr>&& that)
    {
        axis_ = that.axis_;
        matrix_ = std::move(that.matrix_);
        attr_ = that.attr_;
        isOwner_ = that.isOwner_;
        that.isOwner_ = false;
    }

    AdjacencyMatrix(const AdjacencyMatrix<Edge, Attr>& that)
    {
        axis_ = new (std::nothrow) std::vector<u32>();
        *axis_ = *that.axis_;
        matrix_ = that.matrix_;
        attr_ = that.attr_;
        isOwner_ = true;
    }

    AdjacencyMatrix<Edge, Attr>& operator=(AdjacencyMatrix<Edge, Attr>&& that)
    {
        axis_ = that.axis_;
        matrix_ = std::move(that.matrix_);
        attr_ = that.attr_;
        isOwner_ = that.isOwner_;
        that.isOwner_ = false;
        return *this;
    }

    AdjacencyMatrix<Edge, Attr>& operator=(const AdjacencyMatrix<Edge, Attr>& that)
    {
        axis_ = new (std::nothrow) std::vector<u32>();
        *axis_ = *that.axis_;
        matrix_ = that.matrix_;
        attr_ = that.attr_;
        isOwner_ = true;
    }

    // 调用其他方法前必须先执行Init
    HcclResult InitMatrixByList(std::vector<u32> &list)
    {
        axis_ = new (std::nothrow) std::vector<u32>();
        if (axis_ ==  nullptr) {
            return HCCL_E_PTR;
        }
        *axis_ = list;
        matrix_.Init(list.size());

        isOwner_ = true;
        return HCCL_SUCCESS;
    }

    virtual ~AdjacencyMatrix()
    {
        if (isOwner_) {
            delete axis_;
        }
    }

    inline u32 GetNodeNum()
    {
        return axis_->size();
    }

    inline u32 GetNode(u32 index)
    {
        if (index > axis_->size()) {
            return ERR_NODE;
        }
        return (*axis_)[index];
    }

    inline std::vector<u32> GetNodeList()
    {
        return *axis_;
    }

    inline HcclResult AddEdgeByIndex(u32 srcIndex, u32 dstIndex, Edge& edge)
    {
        matrix_(srcIndex, dstIndex) = edge;
        return HCCL_SUCCESS;
    }

    HcclResult RemoveEdgeByIndex(u32 srcIndex, u32 dstIndex)
    {
        if (srcIndex >= axis_->size() || dstIndex >= axis_->size()) {
            return HCCL_E_PARA;
        }
        matrix_(srcIndex, dstIndex) = Edge();
        return HCCL_SUCCESS;
    }

    inline Edge &GetEdgeByIndex(u32 srcIndex, u32 dstIndex)
    {
        return matrix_(srcIndex, dstIndex);
    }

    bool IsConnected(u32 src, u32 dst)
    {
        return matrix_(src, dst).IsConnected();
    }

    inline void SetAttr(Attr attr)
    {
        attr_ = attr;
    }

    inline Attr GetAttr()
    {
        return attr_;
    }
private:
    inline u32 GetIndex(u32 node)
    {
        auto iter = std::find(axis_->begin(), axis_->end(), node);
        if (std::find(axis_->begin(), axis_->end(), node) == axis_->end()) {
            return ERR_NODE;
        }
        return iter - axis_->begin();
    }

    bool isOwner_ = false;
    std::vector<u32> *axis_ = nullptr; // node id
    // std::vector<std::vector<Edge>> *matrix_ = nullptr;
    Matrix<Edge> matrix_;
    Attr attr_;
};

enum class AlgType {
    NOT_SPECIFIED,
    RING,
    MULTI_RING,
    MESH,
    RECURSIVE_HD,
    BINARY_HD,
    PAIR_WISE,
    INVALID_VAL
};

enum class SubgraphOpType {
    ALLTOALL,
    ALLGATHER,
    ALLREDUCE,
    BROADCAST,
    REDUCE_SCATTER,
    GATHER,
    SCATTER,
    REDUCE,
    INVALID_VAL
};

struct SubgraphCalcInfo {
    SubgraphOpType subgraphOpType;
    AlgType   algType;
    u32 root;
    HcclDataType dtype;
    u64 count;
};
} // namespace HcclEvaluator
#endif